import os
import shutil

import cv2
import numpy as np
from SimpleITK import Image
from tqdm import tqdm             #进度条库，可以在 Python 长循环中添加一个进度提示信息。用户只需要封装任意的迭代器，是一个快速、扩展性强的进度条工具库。

import SimpleITK as sitk
import imageio
import matplotlib.pyplot as plt

def slice(seg_path: str, seg_slice_path: str, vol_path:str, vol_slice_path:str):
    id = 0
    for path in os.listdir(vol_path):
        if path.find('mhd') >= 0:
            id += 1
            save_content = os.path.join(seg_slice_path, str(id))
            if os.path.exists(save_content):
                shutil.rmtree(save_content)
            os.makedirs(save_content)

            save_content_2 = os.path.join(vol_slice_path, str(id))
            if os.path.exists(save_content_2):
                shutil.rmtree(save_content_2)
            os.makedirs(save_content_2)

            vol_mhd = sitk.ReadImage(os.path.join(vol_path, path))
            seg_mhd = sitk.ReadImage(os.path.join(seg_path, path))
            (x,y,z) = vol_mhd.GetSize()
            total = x * y

            spacing = vol_mhd.GetSpacing()
            vol = sitk.GetArrayFromImage(vol_mhd)
            seg = sitk.GetArrayFromImage(seg_mhd)

            for i in tqdm(range(len(vol))):
                silce_seg = seg[i]
                if silce_seg.max() == 0:
                    continue
                else:
                    silce_seg = cv2.normalize(silce_seg, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
                    silce_seg = cv2.threshold(silce_seg, 1, 255, cv2.THRESH_BINARY)[1]

                    if (np.sum(silce_seg == 255) / total) > 0.015:
                        silce_vol = vol[i]
                        silce_vol = cv2.normalize(silce_vol, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)

                        # 图像增强：用传统的Unsharp Masking进行锐化处理
                        # 创建模糊图像
                        blurred = cv2.GaussianBlur(silce_vol, (0, 0), 3)
                        # 计算掩模
                        mask = cv2.subtract(silce_vol, blurred)
                        # 应用掩模以增强边缘
                        sharpened = cv2.add(silce_vol, mask)

                        imageio.imwrite(os.path.join(save_content,  f'{id}_{i}.png'), silce_seg)
                        imageio.imwrite(os.path.join(save_content_2,  f'{id}_{i}.png'), sharpened)

seg_path = r'D:\Desktop\graduation project\SLIVER07\data\segmentation'  # 原数据集路径
seg_slice_path = r'D:\Desktop\graduation project\SLIVER07\2d\segmentation'  # 保存路径
vol_path = r"D:\Desktop\graduation project\SLIVER07\data\volume"
vol_slice_path = r"D:\Desktop\graduation project\SLIVER07\2d\volume"

slice(seg_path, seg_slice_path,vol_path,vol_slice_path)